import sys
sys.path.append(r"C:\Users\yanqi\Desktop\rbgame\single\SuperGluePretrainedNetwork")
sys.path.append(r"C:\Users\yanqi\Desktop\rbgame\single\data")
from gn.config import js
# from unh.config import js
from pathlib import Path
import argparse
import cv2
import matplotlib.cm as cm
import torch

from models.matching import Matching
from models.utils import (AverageTimer, VideoStreamer,
                          make_matching_plot_fast, frame2tensor)
import numpy as np
import json
# from img.config import js
import time
from multiprocessing import Pool,Process
import threading
import logging
from logging import handlers
class sgmatch():
    def __init__(self):
        torch.set_grad_enabled(False)
        device = 'cuda'
        config = {
            'superpoint': {
                'nms_radius': 4,
                'keypoint_threshold': 0.005,
                'max_keypoints': -1
            },
            'superglue': {
                'weights': 'indoor',
                'sinkhorn_iterations': 20,
                'match_threshold': 0.2,
            }
        }
        self.matching = Matching(config).eval().to(device)
        keys = ['keypoints', 'scores', 'descriptors']
        img0=cv2.imread(js['templateImageName'],cv2.IMREAD_GRAYSCALE)
        # cv2.imshow('img0',img0)
        # cv2.waitKey(0)
        tensor0 = frame2tensor(img0, device)
        self.last_data = self.matching.superpoint({'image': tensor0})

        ## begin


        # print('last_data0',self.last_data)
        # kp=self.last_data['keypoints'][0].cpu().numpy()
        # sc=self.last_data['scores'][0].cpu().numpy()
        # de=self.last_data['descriptors'][0].cpu().numpy()
        # print(kp.shape,de.shape)
        # imagePoints=[]
        # imageScores=[]
        # imageDescriptors=[]
        # # for index in range(len(kp)):
        # #     imagePoints.append(kp[index])
        # #     imageScores.append(sc[index])
        # #     imageDescriptors.append(de[:,index])
        # img0=cv2.imread(js['templateImageName'])
        # for pos in js['points']:
        #     npos=np.repeat([pos['2D']],len(kp),axis=0)
        #     nnorm=np.linalg.norm(kp-npos,ord=1,axis=1,keepdims=False)
        #     index=nnorm.argmin()
        #     # if nnorm[index]>=10:
        #     #     continue
        #     # print(index)
        #     imagePoints.append(kp[index])
        #     imageScores.append(sc[index])
        #     imageDescriptors.append(de[:,index])
        #     cv2.circle(img0, (round(kp[index][0]), round(kp[index][1])), 1, (0,255,0), 2)
        #     # print(de[:,index].reshape(1,-1)[0])
        # # print(imageDescriptors)
        # cv2.imshow("0",img0)
        # cv2.waitKey(0)
        # imageDescriptors=np.array(imageDescriptors)
        # self.last_data['keypoints']=[torch.tensor(imagePoints).to(device)]
        # self.last_data['scores']=(torch.tensor(imageScores).to(device),)
        # self.last_data['descriptors']=[torch.tensor(imageDescriptors.T).to(device)]
        # print(self.last_data['keypoints'][0].shape,self.last_data['descriptors'][0].shape)
        # print('last_data',self.last_data)

        # ##end
        self.last_data = {k+'0': self.last_data[k] for k in keys}
        self.last_data['image0'] = tensor0
    def match(self,img1):
        # img1=cv2.imread("./img/3.jpg",cv2.IMREAD_GRAYSCALE)
        img1=cv2.cvtColor(img1,cv2.COLOR_BGR2GRAY)
        tensor1 = frame2tensor(img1, 'cuda')
        ti=time.time()
        pred = self.matching({**self.last_data, 'image1': tensor1})
        kpts0 = self.last_data['keypoints0'][0].cpu().numpy()
        kpts1 = pred['keypoints1'][0].cpu().numpy()
        matches = pred['matches0'][0].cpu().numpy()
        confidence = pred['matching_scores0'][0].cpu().numpy()
        print(f'ti:{time.time()-ti}')

        valid = matches > -1
        mkpts0 = kpts0[valid]
        mkpts1 = kpts1[matches[valid]]
        imagePoints=[]
        objectPoints=[]
        img0=cv2.imread(js['templateImageName'])
        # for index in range(len(mkpts0)):
        #     cv2.circle(img0, (round(mkpts0[index][0]), round(mkpts0[index][1])), 1, (0,255,0), 2)
        # cv2.imshow("img0",img0)
        
        for pos in js['points']:
            npos=np.repeat([pos['2D']],len(mkpts0),axis=0)
            nnorm=np.linalg.norm(np.array(mkpts0)-npos,ord=1,axis=1,keepdims=False)
            index=nnorm.argmin()
            if nnorm[index]>=10:
                continue
            print('result',mkpts0[index],mkpts1[index],confidence[index])
            # if confidence[index]<=0.5:
            #     continue
            # cv2.circle(img0, (round(mkpts0[index][0]), round(mkpts0[index][1])), 1, (0,0,255), 2)
            # cv2.circle(img1, (round(mkpts1[index][0]), round(mkpts1[index][1])), 1, 255, 2)
            imagePoints.append((mkpts1[index][0],mkpts1[index][1]))
            objectPoints.append(pos['3D'])
        # cv2.imshow("img0",img0)
        return imagePoints,objectPoints
if __name__ == '__main__':
    s=sgmatch()
    img1=cv2.imread(r"C:\Users\yanqi\Desktop\rbgame\single\SuperGluePretrainedNetwork\img\21.jpg")
    print(s.match(img1))
    
